"""
Demo3: 象棋棋子识别 - 实战应用
===============================

这个demo将教你：
1. 如何使用OpenCV处理图像
2. 如何进行图像预处理和特征提取
3. 如何使用迁移学习
4. 如何构建实际应用场景的模型
5. 如何创建简单的图像识别系统

注意：这个demo会创建一个简化的象棋棋子识别系统
由于没有真实的象棋数据集，我们会使用模拟数据来演示流程
"""

import numpy as np
import matplotlib.pyplot as plt
import cv2
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import os
import warnings
warnings.filterwarnings('ignore')

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

print("=" * 60)
print("欢迎来到实战应用 - 象棋棋子识别")
print("=" * 60)

# ============================================
# 步骤1: 创建模拟数据集
# ============================================
print("\n【步骤1】创建模拟象棋棋子数据集...")
print("注意：这是教学演示，使用模拟数据")

# 定义象棋棋子类别
chess_pieces = {
    0: '车', 1: '马', 2: '象', 3: '士', 4: '将',
    5: '炮', 6: '兵'
}

def create_simulated_chess_piece(piece_type, size=64):
    """创建模拟的象棋棋子图像"""
    img = np.ones((size, size, 3), dtype=np.uint8) * 255
    
    # 绘制圆形背景
    center = (size // 2, size // 2)
    radius = size // 3
    
    # 不同棋子用不同颜色和形状
    colors = [
        (255, 100, 100),  # 车 - 红色
        (100, 255, 100),  # 马 - 绿色
        (100, 100, 255),  # 象 - 蓝色
        (255, 255, 100),  # 士 - 黄色
        (255, 100, 255),  # 将 - 紫色
        (100, 255, 255),  # 炮 - 青色
        (200, 200, 100),  # 兵 - 橄榄色
    ]
    
    color = colors[piece_type]
    cv2.circle(img, center, radius, color, -1)
    cv2.circle(img, center, radius, (0, 0, 0), 2)
    
    # 添加一些随机噪声使数据更真实
    noise = np.random.randint(-30, 30, img.shape, dtype=np.int16)
    img = np.clip(img.astype(np.int16) + noise, 0, 255).astype(np.uint8)
    
    # 随机旋转
    angle = np.random.randint(-15, 15)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    img = cv2.warpAffine(img, M, (size, size))
    
    return img

# 创建数据集
print("生成训练数据...")
num_samples_per_class = 200
X_train_list = []
y_train_list = []

for piece_type in range(7):
    for _ in range(num_samples_per_class):
        img = create_simulated_chess_piece(piece_type)
        X_train_list.append(img)
        y_train_list.append(piece_type)

X_train = np.array(X_train_list)
y_train = np.array(y_train_list)

# 创建测试集
print("生成测试数据...")
num_test_samples = 30
X_test_list = []
y_test_list = []

for piece_type in range(7):
    for _ in range(num_test_samples):
        img = create_simulated_chess_piece(piece_type)
        X_test_list.append(img)
        y_test_list.append(piece_type)

X_test = np.array(X_test_list)
y_test = np.array(y_test_list)

print(f"✓ 数据集创建完成")
print(f"  - 训练集样本数: {len(X_train)}")
print(f"  - 测试集样本数: {len(X_test)}")
print(f"  - 图片尺寸: {X_train.shape[1]} x {X_train.shape[2]}")
print(f"  - 棋子类别: {len(chess_pieces)}")

# ============================================
# 步骤2: 数据可视化
# ============================================
print("\n【步骤2】数据可视化...")

fig, axes = plt.subplots(2, 7, figsize=(15, 5))
fig.suptitle('模拟象棋棋子样本', fontsize=16, fontweight='bold')

for i in range(7):
    # 显示两个样本
    for j in range(2):
        ax = axes[j, i]
        sample_idx = i * num_samples_per_class + j
        ax.imshow(X_train[sample_idx])
        if j == 0:
            ax.set_title(chess_pieces[i], fontsize=12, fontweight='bold')
        ax.axis('off')

plt.tight_layout()
plt.savefig('chess_samples.png', dpi=300, bbox_inches='tight')
print("✓ 样本图片已保存: chess_samples.png")

# ============================================
# 步骤3: 图像预处理
# ============================================
print("\n【步骤3】图像预处理...")

# 归一化
X_train = X_train.astype('float32') / 255.0
X_test = X_test.astype('float32') / 255.0

# 标签one-hot编码
y_train_cat = keras.utils.to_categorical(y_train, 7)
y_test_cat = keras.utils.to_categorical(y_test, 7)

print(f"✓ 数据预处理完成")
print(f"  - 训练集形状: {X_train.shape}")
print(f"  - 测试集形状: {X_test.shape}")

# ============================================
# 步骤4: 数据增强
# ============================================
print("\n【步骤4】配置数据增强...")

datagen = ImageDataGenerator(
    rotation_range=20,
    width_shift_range=0.1,
    height_shift_range=0.1,
    zoom_range=0.1,
    horizontal_flip=False,
    fill_mode='nearest'
)

print("✓ 数据增强配置完成")
print("  - 旋转范围: ±20度")
print("  - 平移范围: ±10%")
print("  - 缩放范围: ±10%")

# 展示数据增强效果
sample_img = X_train[0:1]
fig, axes = plt.subplots(2, 5, figsize=(12, 5))
fig.suptitle('数据增强效果展示', fontsize=16, fontweight='bold')

axes[0, 0].imshow(sample_img[0])
axes[0, 0].set_title('原始图像', fontweight='bold')
axes[0, 0].axis('off')

aug_iter = datagen.flow(sample_img, batch_size=1)
for i in range(1, 10):
    ax = axes[i // 5, i % 5]
    aug_img = next(aug_iter)[0]
    ax.imshow(aug_img)
    ax.set_title(f'增强 {i}')
    ax.axis('off')

plt.tight_layout()
plt.savefig('data_augmentation.png', dpi=300, bbox_inches='tight')
print("✓ 数据增强效果图已保存: data_augmentation.png")

# ============================================
# 步骤5: 构建自定义CNN模型
# ============================================
print("\n【步骤5】构建自定义CNN模型...")

model_custom = keras.Sequential([
    layers.Input(shape=(64, 64, 3)),
    
    # 第一个卷积块
    layers.Conv2D(32, (3, 3), activation='relu', padding='same'),
    layers.BatchNormalization(),
    layers.Conv2D(32, (3, 3), activation='relu', padding='same'),
    layers.MaxPooling2D((2, 2)),
    layers.Dropout(0.25),
    
    # 第二个卷积块
    layers.Conv2D(64, (3, 3), activation='relu', padding='same'),
    layers.BatchNormalization(),
    layers.Conv2D(64, (3, 3), activation='relu', padding='same'),
    layers.MaxPooling2D((2, 2)),
    layers.Dropout(0.25),
    
    # 第三个卷积块
    layers.Conv2D(128, (3, 3), activation='relu', padding='same'),
    layers.BatchNormalization(),
    layers.Conv2D(128, (3, 3), activation='relu', padding='same'),
    layers.MaxPooling2D((2, 2)),
    layers.Dropout(0.25),
    
    # 全连接层
    layers.Flatten(),
    layers.Dense(256, activation='relu'),
    layers.BatchNormalization(),
    layers.Dropout(0.5),
    layers.Dense(7, activation='softmax')
], name='Custom_CNN')

model_custom.compile(
    optimizer='adam',
    loss='categorical_crossentropy',
    metrics=['accuracy']
)

print("✓ 自定义CNN模型构建完成")
print("\n模型结构:")
model_custom.summary()

# ============================================
# 步骤6: 训练自定义模型
# ============================================
print("\n【步骤6】训练自定义CNN模型...")
print("训练中，请稍候...")

history_custom = model_custom.fit(
    datagen.flow(X_train, y_train_cat, batch_size=32),
    epochs=20,
    validation_data=(X_test, y_test_cat),
    verbose=1
)

print("✓ 自定义CNN模型训练完成")

# ============================================
# 步骤7: 构建迁移学习模型
# ============================================
print("\n【步骤7】构建迁移学习模型（使用MobileNetV2）...")

# 加载预训练的MobileNetV2
base_model = MobileNetV2(
    input_shape=(64, 64, 3),
    include_top=False,
    weights='imagenet'
)

# 冻结基础模型
base_model.trainable = False

model_transfer = keras.Sequential([
    base_model,
    layers.GlobalAveragePooling2D(),
    layers.Dense(128, activation='relu'),
    layers.Dropout(0.5),
    layers.Dense(7, activation='softmax')
], name='Transfer_Learning')

model_transfer.compile(
    optimizer='adam',
    loss='categorical_crossentropy',
    metrics=['accuracy']
)

print("✓ 迁移学习模型构建完成")
print("\n模型结构:")
model_transfer.summary()

# ============================================
# 步骤8: 训练迁移学习模型
# ============================================
print("\n【步骤8】训练迁移学习模型...")
print("训练中，请稍候...")

history_transfer = model_transfer.fit(
    X_train, y_train_cat,
    batch_size=32,
    epochs=10,
    validation_data=(X_test, y_test_cat),
    verbose=1
)

print("✓ 迁移学习模型训练完成")

# ============================================
# 步骤9: 模型评估和对比
# ============================================
print("\n【步骤9】模型评估和对比...")

# 评估自定义CNN
loss_custom, acc_custom = model_custom.evaluate(X_test, y_test_cat, verbose=0)
print(f"\n自定义CNN模型:")
print(f"  - 测试集损失: {loss_custom:.4f}")
print(f"  - 测试集准确率: {acc_custom:.4f} ({acc_custom*100:.2f}%)")

# 评估迁移学习模型
loss_transfer, acc_transfer = model_transfer.evaluate(X_test, y_test_cat, verbose=0)
print(f"\n迁移学习模型:")
print(f"  - 测试集损失: {loss_transfer:.4f}")
print(f"  - 测试集准确率: {acc_transfer:.4f} ({acc_transfer*100:.2f}%)")

# 确定最佳模型
if acc_transfer > acc_custom:
    best_model_name = "迁移学习模型"
    best_model = model_transfer
    best_acc = acc_transfer
else:
    best_model_name = "自定义CNN"
    best_model = model_custom
    best_acc = acc_custom

print(f"\n{'='*60}")
print(f"🏆 最佳模型: {best_model_name}")
print(f"🎯 准确率: {best_acc:.4f} ({best_acc*100:.2f}%)")
print(f"{'='*60}")

# ============================================
# 步骤10: 可视化训练过程
# ============================================
print("\n【步骤10】可视化训练过程...")

fig, axes = plt.subplots(2, 2, figsize=(14, 10))
fig.suptitle('模型训练过程对比', fontsize=16, fontweight='bold')

# 自定义CNN - 准确率
ax1 = axes[0, 0]
ax1.plot(history_custom.history['accuracy'], label='训练准确率', linewidth=2)
ax1.plot(history_custom.history['val_accuracy'], label='验证准确率', linewidth=2)
ax1.set_xlabel('Epoch')
ax1.set_ylabel('准确率')
ax1.set_title('自定义CNN - 准确率')
ax1.legend()
ax1.grid(True, alpha=0.3)

# 自定义CNN - 损失
ax2 = axes[0, 1]
ax2.plot(history_custom.history['loss'], label='训练损失', linewidth=2)
ax2.plot(history_custom.history['val_loss'], label='验证损失', linewidth=2)
ax2.set_xlabel('Epoch')
ax2.set_ylabel('损失')
ax2.set_title('自定义CNN - 损失')
ax2.legend()
ax2.grid(True, alpha=0.3)

# 迁移学习 - 准确率
ax3 = axes[1, 0]
ax3.plot(history_transfer.history['accuracy'], label='训练准确率', linewidth=2)
ax3.plot(history_transfer.history['val_accuracy'], label='验证准确率', linewidth=2)
ax3.set_xlabel('Epoch')
ax3.set_ylabel('准确率')
ax3.set_title('迁移学习模型 - 准确率')
ax3.legend()
ax3.grid(True, alpha=0.3)

# 迁移学习 - 损失
ax4 = axes[1, 1]
ax4.plot(history_transfer.history['loss'], label='训练损失', linewidth=2)
ax4.plot(history_transfer.history['val_loss'], label='验证损失', linewidth=2)
ax4.set_xlabel('Epoch')
ax4.set_ylabel('损失')
ax4.set_title('迁移学习模型 - 损失')
ax4.legend()
ax4.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('training_comparison.png', dpi=300, bbox_inches='tight')
print("✓ 训练对比图已保存: training_comparison.png")

# ============================================
# 步骤11: 预测示例
# ============================================
print("\n【步骤11】预测示例...")

# 选择一些测试样本
num_samples = 14
sample_indices = np.random.choice(len(X_test), num_samples, replace=False)

fig, axes = plt.subplots(2, 7, figsize=(15, 5))
fig.suptitle(f'{best_model_name} - 预测结果', fontsize=16, fontweight='bold')

for i, idx in enumerate(sample_indices):
    ax = axes[i // 7, i % 7]
    
    # 显示图片
    ax.imshow(X_test[idx])
    
    # 预测
    pred = best_model.predict(X_test[idx:idx+1], verbose=0)
    pred_class = np.argmax(pred)
    true_class = y_test[idx]
    confidence = np.max(pred) * 100
    
    # 设置标题
    color = 'green' if pred_class == true_class else 'red'
    ax.set_title(f'真: {chess_pieces[true_class]}\n'
                 f'预测: {chess_pieces[pred_class]}\n'
                 f'{confidence:.1f}%',
                 color=color, fontsize=9, fontweight='bold')
    ax.axis('off')

plt.tight_layout()
plt.savefig('predictions.png', dpi=300, bbox_inches='tight')
print("✓ 预测结果已保存: predictions.png")

# ============================================
# 步骤12: 保存模型
# ============================================
print("\n【步骤12】保存模型...")

model_custom.save('custom_cnn_model.h5')
model_transfer.save('transfer_learning_model.h5')

print("✓ 模型已保存:")
print("  - custom_cnn_model.h5")
print("  - transfer_learning_model.h5")

# ============================================
# 步骤13: 创建简单的识别函数
# ============================================
print("\n【步骤13】创建识别函数...")

def recognize_chess_piece(image, model):
    """识别象棋棋子"""
    # 预处理
    if image.shape != (64, 64, 3):
        image = cv2.resize(image, (64, 64))
    image = image.astype('float32') / 255.0
    image = np.expand_dims(image, axis=0)
    
    # 预测
    pred = model.predict(image, verbose=0)
    pred_class = np.argmax(pred)
    confidence = np.max(pred)
    
    return chess_pieces[pred_class], confidence

# 测试识别函数
test_img = create_simulated_chess_piece(2)  # 创建一个"象"
piece_name, confidence = recognize_chess_piece(test_img, best_model)

print(f"✓ 识别函数创建完成")
print(f"\n测试识别:")
print(f"  - 识别结果: {piece_name}")
print(f"  - 置信度: {confidence*100:.2f}%")

# ============================================
# 总结
# ============================================
print("\n" + "="*60)
print("🎉 恭喜！你已经完成了所有机器学习入门项目！")
print("="*60)
print("\n在这个项目中，你学到了：")
print("  ✓ 如何使用OpenCV处理图像")
print("  ✓ 如何创建和使用数据增强")
print("  ✓ 如何构建复杂的CNN模型")
print("  ✓ 如何使用迁移学习")
print("  ✓ 如何对比不同模型的性能")
print("  ✓ 如何创建实际应用的识别系统")

print("\n💡 实战扩展建议：")
print("  1. 收集真实的象棋棋子图片进行训练")
print("  2. 添加棋盘检测功能")
print("  3. 实现完整的棋局识别系统")
print("  4. 添加实时摄像头识别功能")
print("  5. 结合象棋AI引擎，实现智能提示")

print("\n🚀 下一步学习方向：")
print("  1. 深入学习计算机视觉（目标检测、图像分割）")
print("  2. 学习自然语言处理（NLP）")
print("  3. 参加Kaggle竞赛")
print("  4. 开发自己的机器学习项目")
print("  5. 学习模型部署（Flask、FastAPI、TensorFlow Serving）")

print("\n📚 推荐资源：")
print("  - OpenCV官方教程: https://docs.opencv.org/")
print("  - TensorFlow官方教程: https://www.tensorflow.org/tutorials")
print("  - Kaggle竞赛平台: https://www.kaggle.com/")
print("  - Papers with Code: https://paperswithcode.com/")

print("\n" + "="*60)
print("感谢你完成这个学习旅程！")
print("继续保持学习热情，不断实践和探索！")
print("="*60)

plt.show()
